import matplotlib.pyplot as plt


def plot_hist(data_df, ppc_data=None, bins=30):
    # Plotting the RTs

    if ppc_data is None:
        plt.hist(
            data_df["rt"] * data_df["congruency"],
            histtype="step",
            color="black",
            density=True,
            bins=bins,
        )
    else:
        plt.hist(
            data_df["rt"] * data_df["congruency"],
            histtype="step",
            color="black",
            density=True,
            bins=bins,
            label="Original Data",
        )
        plt.hist(
            ppc_data["rt"] * ppc_data["congruency"],
            histtype="step",
            color="red",
            density=True,
            bins=bins,
            label="Posterior Mean",
        )
        plt.legend()
    plt.xlabel("Reaction Time")
    plt.ylabel("Density")
    plt.show()




def _sum_activation(trajectories):
    # trajectories = [[1, 2, 3], [4, 5], [6, 7, 8, 9]]
    # 初始化sum_activation列表
    sum_activation = []
    # 使用zip_longest来遍历trajectories中的每个trajectory
    for traj in trajectories:
        # 如果sum_activation的长度小于当前trajectory的长度，则用0填充
        while len(sum_activation) < len(traj):
            sum_activation.append(0)
    # 累加当前trajectory的值到sum_activation
        for i in range(len(traj)):
            sum_activation[i] += traj[i]
    print(sum_activation)


def plot_DDM(ntrials: int, diffusion_trial:callable, dt=0.001,**kwargs):

    from matplotlib.gridspec import GridSpec
    import seaborn as sns
    import pandas as pd
    import numpy as np

    fig = plt.figure()
    gs = GridSpec(3, 1, height_ratios=[1, 3, 1], hspace=0)
    ax1 = plt.subplot(gs[0])
    ax2 = plt.subplot(gs[1])
    ax3 = plt.subplot(gs[2])
    rts = []
    choices = []
    counter = 0
    tmp_v_view = []
    max_evidence = 0
    sum_activation = []
    while counter < ntrials:
        rt, choice, v_view, trajectory = diffusion_trial(**kwargs)
        if choice is not np.nan:
            rts.append(rt)
            choices.append(choice)
            if len(v_view) > len(tmp_v_view):
                tmp_v_view = v_view
            max_evidence = max(max_evidence, np.max(trajectory))
            max_evidence_len = len(trajectory)
            sum_activation += [trajectory]
            # if choice == 1:
            #     while len(sum_activation) < len(trajectory):
            #         sum_activation.append(0)
            #     for i in range(len(trajectory)):
            #         sum_activation[i] += trajectory[i]
            counter += 1

        ax2.plot(
            np.arange(len(trajectory)) * dt, trajectory, color="gray", alpha=0.3
        )

    x = np.arange(2000)
    # the mean evidence for control drift 
    vc = tmp_v_view[0]
    # vc = kwargs["theta"][2]
    control_activation = vc * x  * dt + trajectory[0]
    control_activation = control_activation[control_activation <= max_evidence]
    control_activation = control_activation[:max_evidence_len]
    ax2.plot(np.arange(len(control_activation)) * dt, control_activation, color="blue", label="control")
    # sum_activation = to_array(sum_activation)
    # sum_activation = np.nanmean(sum_activation, axis=0)
    # sum_activation = sum_activation[sum_activation <= max_evidence]
    # ax2.plot(np.arange(len(sum_activation)) * dt, sum_activation, color="red")

    # the mean evidence for varied drift 
    tmp_v_view2 = np.cumsum(tmp_v_view)*dt + trajectory[0]
    # tmp_v_view2 = np.array(tmp_v_view)*10*np.arange(len(tmp_v_view)) + trajectory[0]
    tmp_v_view2 = tmp_v_view2[tmp_v_view2 <= max_evidence]
    tmp_v_view2 = tmp_v_view2[:max_evidence_len]
    ax2.plot(np.arange(len(tmp_v_view2)) * dt, tmp_v_view2, color="yellow")
    # 设置标题和横纵坐标标签
    ax2.set_xlabel("RT(seconds)")
    ax2.set_ylabel('evidence')

    aa = pd.DataFrame({"rt": rts, "choice": choices})

    # 绘制上方的密度分布图
    sns.histplot(
        aa[aa.choice == 1].rt, ax=ax1, color='orange', linewidth=1.5, kde=True
    )

    # 绘制下方的倒置密度分布图
    sns.histplot(
        aa[aa.choice != 1].rt, ax=ax3, color='blue', linewidth=1.5, kde=True
    )

    # 获取两个轴的最大计数
    max_count = max(ax1.get_ylim()[1], ax3.get_ylim()[1])
    # 设置两个轴的y轴限制相同
    ax1.set_ylim(0, max_count)
    ax3.set_ylim(0, max_count)

    ax1.set_axis_off()
    ax3.set_axis_off()
    ax3.invert_yaxis()  # 倒置y轴
    # 调整图形布局
    fig.subplots_adjust(hspace=0.01)  # 调整子图之间的间距

    return fig